"""
券商管理器

管理多个券商连接，提供统一的接口。
"""

import asyncio
import logging
from typing import Dict, List, Optional, Any, Callable

from .base_broker import BaseBroker, BrokerType, Order, OrderRequest, Trade, Position, AccountInfo, MarketData
from .mock_broker import MockBroker


class BrokerManager:
    """券商管理器"""
    
    def __init__(self):
        self.logger = logging.getLogger(__name__)
        
        # 券商实例
        self.brokers: Dict[str, BaseBroker] = {}
        
        # 默认券商
        self.default_broker_id: Optional[str] = None
        
        # 回调函数
        self.order_update_callback: Optional[Callable[[str, Order], None]] = None
        self.trade_callback: Optional[Callable[[str, Trade], None]] = None
        self.market_data_callback: Optional[Callable[[str, MarketData], None]] = None
        
        # 内部状态
        self._lock = asyncio.Lock()
    
    async def add_broker(self, broker_id: str, broker_type: BrokerType, 
                        config: Dict[str, Any]) -> bool:
        """添加券商"""
        try:
            # 创建券商实例
            if broker_type == BrokerType.MOCK:
                broker = MockBroker(config)
            else:
                # 其他券商类型的实现
                raise NotImplementedError(f"券商类型 {broker_type.value} 尚未实现")
            
            # 设置回调函数
            broker.set_order_update_callback(
                lambda order: self._on_order_update(broker_id, order)
            )
            broker.set_trade_callback(
                lambda trade: self._on_trade(broker_id, trade)
            )
            broker.set_market_data_callback(
                lambda market_data: self._on_market_data(broker_id, market_data)
            )
            
            # 连接券商
            if await broker.connect():
                async with self._lock:
                    self.brokers[broker_id] = broker
                    
                    # 如果是第一个券商，设为默认券商
                    if self.default_broker_id is None:
                        self.default_broker_id = broker_id
                
                self.logger.info(f"券商 {broker_id} 添加成功")
                return True
            else:
                self.logger.error(f"券商 {broker_id} 连接失败")
                return False
                
        except Exception as e:
            self.logger.error(f"添加券商 {broker_id} 失败: {e}")
            return False
    
    async def remove_broker(self, broker_id: str) -> bool:
        """移除券商"""
        async with self._lock:
            broker = self.brokers.get(broker_id)
            if broker:
                await broker.disconnect()
                del self.brokers[broker_id]
                
                # 如果移除的是默认券商，重新选择默认券商
                if self.default_broker_id == broker_id:
                    self.default_broker_id = next(iter(self.brokers.keys()), None)
                
                self.logger.info(f"券商 {broker_id} 已移除")
                return True
            
            return False
    
    async def get_broker(self, broker_id: Optional[str] = None) -> Optional[BaseBroker]:
        """获取券商实例"""
        if broker_id is None:
            broker_id = self.default_broker_id
        
        return self.brokers.get(broker_id) if broker_id else None
    
    async def set_default_broker(self, broker_id: str) -> bool:
        """设置默认券商"""
        if broker_id in self.brokers:
            self.default_broker_id = broker_id
            self.logger.info(f"默认券商设置为: {broker_id}")
            return True
        return False
    
    async def submit_order(self, order_request: OrderRequest, 
                          broker_id: Optional[str] = None) -> Optional[Order]:
        """提交订单"""
        broker = await self.get_broker(broker_id)
        if broker:
            return await broker.submit_order(order_request)
        return None
    
    async def cancel_order(self, order_id: str, 
                          broker_id: Optional[str] = None) -> bool:
        """撤销订单"""
        broker = await self.get_broker(broker_id)
        if broker:
            return await broker.cancel_order(order_id)
        return False
    
    async def get_order(self, order_id: str, 
                       broker_id: Optional[str] = None) -> Optional[Order]:
        """获取订单信息"""
        broker = await self.get_broker(broker_id)
        if broker:
            return await broker.get_order(order_id)
        return None
    
    async def get_orders(self, symbol: Optional[str] = None, 
                        broker_id: Optional[str] = None) -> List[Order]:
        """获取订单列表"""
        broker = await self.get_broker(broker_id)
        if broker:
            return await broker.get_orders(symbol)
        return []
    
    async def get_all_orders(self, symbol: Optional[str] = None) -> Dict[str, List[Order]]:
        """获取所有券商的订单列表"""
        all_orders = {}
        for broker_id, broker in self.brokers.items():
            try:
                orders = await broker.get_orders(symbol)
                all_orders[broker_id] = orders
            except Exception as e:
                self.logger.error(f"获取券商 {broker_id} 订单失败: {e}")
                all_orders[broker_id] = []
        
        return all_orders
    
    async def get_trades(self, symbol: Optional[str] = None, 
                        broker_id: Optional[str] = None) -> List[Trade]:
        """获取成交记录"""
        broker = await self.get_broker(broker_id)
        if broker:
            return await broker.get_trades(symbol)
        return []
    
    async def get_all_trades(self, symbol: Optional[str] = None) -> Dict[str, List[Trade]]:
        """获取所有券商的成交记录"""
        all_trades = {}
        for broker_id, broker in self.brokers.items():
            try:
                trades = await broker.get_trades(symbol)
                all_trades[broker_id] = trades
            except Exception as e:
                self.logger.error(f"获取券商 {broker_id} 成交记录失败: {e}")
                all_trades[broker_id] = []
        
        return all_trades
    
    async def get_positions(self, broker_id: Optional[str] = None) -> List[Position]:
        """获取持仓信息"""
        broker = await self.get_broker(broker_id)
        if broker:
            return await broker.get_positions()
        return []
    
    async def get_all_positions(self) -> Dict[str, List[Position]]:
        """获取所有券商的持仓信息"""
        all_positions = {}
        for broker_id, broker in self.brokers.items():
            try:
                positions = await broker.get_positions()
                all_positions[broker_id] = positions
            except Exception as e:
                self.logger.error(f"获取券商 {broker_id} 持仓信息失败: {e}")
                all_positions[broker_id] = []
        
        return all_positions
    
    async def get_account_info(self, broker_id: Optional[str] = None) -> Optional[AccountInfo]:
        """获取账户信息"""
        broker = await self.get_broker(broker_id)
        if broker:
            return await broker.get_account_info()
        return None
    
    async def get_all_account_info(self) -> Dict[str, Optional[AccountInfo]]:
        """获取所有券商的账户信息"""
        all_accounts = {}
        for broker_id, broker in self.brokers.items():
            try:
                account_info = await broker.get_account_info()
                all_accounts[broker_id] = account_info
            except Exception as e:
                self.logger.error(f"获取券商 {broker_id} 账户信息失败: {e}")
                all_accounts[broker_id] = None
        
        return all_accounts
    
    async def get_market_data(self, symbol: str, 
                             broker_id: Optional[str] = None) -> Optional[MarketData]:
        """获取市场数据"""
        broker = await self.get_broker(broker_id)
        if broker:
            return await broker.get_market_data(symbol)
        return None
    
    async def subscribe_market_data(self, symbols: List[str], 
                                   broker_id: Optional[str] = None) -> bool:
        """订阅市场数据"""
        broker = await self.get_broker(broker_id)
        if broker:
            return await broker.subscribe_market_data(symbols)
        return False
    
    async def unsubscribe_market_data(self, symbols: List[str], 
                                     broker_id: Optional[str] = None) -> bool:
        """取消订阅市场数据"""
        broker = await self.get_broker(broker_id)
        if broker:
            return await broker.unsubscribe_market_data(symbols)
        return False
    
    async def subscribe_all_market_data(self, symbols: List[str]) -> Dict[str, bool]:
        """在所有券商上订阅市场数据"""
        results = {}
        for broker_id, broker in self.brokers.items():
            try:
                result = await broker.subscribe_market_data(symbols)
                results[broker_id] = result
            except Exception as e:
                self.logger.error(f"券商 {broker_id} 订阅市场数据失败: {e}")
                results[broker_id] = False
        
        return results
    
    def set_order_update_callback(self, callback: Callable[[str, Order], None]):
        """设置订单更新回调"""
        self.order_update_callback = callback
    
    def set_trade_callback(self, callback: Callable[[str, Trade], None]):
        """设置成交回调"""
        self.trade_callback = callback
    
    def set_market_data_callback(self, callback: Callable[[str, MarketData], None]):
        """设置市场数据回调"""
        self.market_data_callback = callback
    
    def _on_order_update(self, broker_id: str, order: Order):
        """订单更新回调"""
        if self.order_update_callback:
            try:
                self.order_update_callback(broker_id, order)
            except Exception as e:
                self.logger.error(f"订单更新回调错误: {e}")
    
    def _on_trade(self, broker_id: str, trade: Trade):
        """成交回调"""
        if self.trade_callback:
            try:
                self.trade_callback(broker_id, trade)
            except Exception as e:
                self.logger.error(f"成交回调错误: {e}")
    
    def _on_market_data(self, broker_id: str, market_data: MarketData):
        """市场数据回调"""
        if self.market_data_callback:
            try:
                self.market_data_callback(broker_id, market_data)
            except Exception as e:
                self.logger.error(f"市场数据回调错误: {e}")
    
    async def get_broker_list(self) -> List[Dict[str, Any]]:
        """获取券商列表"""
        broker_list = []
        for broker_id, broker in self.brokers.items():
            broker_info = broker.get_broker_info()
            broker_info['broker_id'] = broker_id
            broker_info['is_default'] = broker_id == self.default_broker_id
            broker_list.append(broker_info)
        
        return broker_list
    
    async def get_statistics(self) -> Dict[str, Any]:
        """获取统计信息"""
        stats = {
            'total_brokers': len(self.brokers),
            'connected_brokers': sum(1 for broker in self.brokers.values() if broker.is_connected),
            'default_broker': self.default_broker_id,
            'broker_types': {}
        }
        
        # 统计券商类型
        for broker in self.brokers.values():
            broker_type = broker.broker_type.value
            stats['broker_types'][broker_type] = stats['broker_types'].get(broker_type, 0) + 1
        
        return stats
    
    async def disconnect_all(self):
        """断开所有券商连接"""
        for broker_id, broker in self.brokers.items():
            try:
                await broker.disconnect()
                self.logger.info(f"券商 {broker_id} 已断开连接")
            except Exception as e:
                self.logger.error(f"断开券商 {broker_id} 连接失败: {e}")
        
        async with self._lock:
            self.brokers.clear()
            self.default_broker_id = None
        
        self.logger.info("所有券商连接已断开")